#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：TU1.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/10/23 15:15 
'''
# !/usr/bin/env python
# -*- coding: UTF-8 -*-
import os
import json
import gurobipy as gp
from gurobipy import GRB
import numpy as np
from typing import Dict, List


class GasSchedulingGurobiSolver:
    def __init__(self, instance_path: str):
        self.instance = self._load_instance(instance_path)
        self._parse_parameters()
        self.model = None
        self._init_model()

    def _load_instance(self, path: str) -> Dict:
        if not os.path.exists(path):
            raise FileNotFoundError(f"算例文件不存在: {path}")
        with open(path, 'r', encoding='utf-8') as f:
            return json.load(f)

    def _parse_parameters(self):
        params = self.instance["parameters"]
        times = self.instance["processing_times"]

        # 基本参数
        self.num_workpieces = params["num_workpieces"]  # 工件数
        self.num_components = params["num_components"]  # 组分数
        self.num_inflation_eq = params["num_inflation_eq"]  # 充装设备数
        self.num_analysis_eq_per_component = params["num_analysis_eq_per_component"]
        self.num_analysis_eq = self.num_components * self.num_analysis_eq_per_component  # 总分析设备数
        self.num_clean_eq = params["num_clean_eq"]  # 清洗设备数
        self.num_shake_eq = params["num_shake_eq"]  # 摇匀设备数

        # 时间参数
        self.clean_time = times["clean_time"]
        self.shake_unit_time = times["shake_unit_time"]
        self.inflation_time = np.array(times["inflation_time"])  # 每个工件的总充装时间（后续可拆分为组分级）
        self.analysis_time = np.array(times["analysis_time"])  # [工件i, 组分c]的分析时间

        # 分析设备能力（修正：确保每个组分至少被1台设备覆盖，且设备覆盖2个随机组分，避免漏洞）
        self.analysis_eq_capabilities = self._generate_analysis_capabilities()

    def _generate_analysis_capabilities(self) -> List[List[int]]:
        """生成更合理的设备能力：每个设备覆盖1-2个组分，确保所有组分被覆盖"""
        capabilities = []
        # 先确保每个组分有至少1台设备单独覆盖（避免串联依赖）
        for comp in range(self.num_components):
            capabilities.append([comp])  # 设备comp：仅覆盖组分comp
        # 剩余设备随机覆盖2个不同组分（增加灵活性）
        remaining_eq = self.num_analysis_eq - self.num_components
        for _ in range(remaining_eq):
            if self.num_components >= 2:
                # 随机选2个不同组分
                c1, c2 = np.random.choice(self.num_components, 2, replace=False)
                capabilities.append([c1, c2])
            else:
                # 只有1个组分时，设备只能覆盖它
                capabilities.append([0])
        return capabilities

    def _init_model(self):
        self.model = gp.Model("GasScheduling")
        self.model.setParam("OutputFlag", 1)  # 显示求解过程
        self.model.setParam("NumericFocus", 3)  # 增强数值稳定性

    def build_model(self):
        # -------------------------- 1. 决策变量 --------------------------
        # 清洗阶段：工件i在设备c上的开始/结束时间及使用标记
        clean_start = self.model.addVars(
            self.num_workpieces, self.num_clean_eq,
            vtype=GRB.CONTINUOUS, name="clean_start"
        )
        clean_end = self.model.addVars(
            self.num_workpieces, self.num_clean_eq,
            vtype=GRB.CONTINUOUS, name="clean_end"
        )
        clean_use = self.model.addVars(
            self.num_workpieces, self.num_clean_eq,
            vtype=GRB.BINARY, name="clean_use"
        )

        # 充装阶段：工件i在设备f上的开始/结束时间及使用标记
        inflate_start = self.model.addVars(
            self.num_workpieces, self.num_inflation_eq,
            vtype=GRB.CONTINUOUS, name="inflate_start"
        )
        inflate_end = self.model.addVars(
            self.num_workpieces, self.num_inflation_eq,
            vtype=GRB.CONTINUOUS, name="inflate_end"
        )
        inflate_use = self.model.addVars(
            self.num_workpieces, self.num_inflation_eq,
            vtype=GRB.BINARY, name="inflate_use"
        )

        # 摇匀阶段：工件i在设备s上的开始/结束时间及使用标记
        shake_start = self.model.addVars(
            self.num_workpieces, self.num_shake_eq,
            vtype=GRB.CONTINUOUS, name="shake_start"
        )
        shake_end = self.model.addVars(
            self.num_workpieces, self.num_shake_eq,
            vtype=GRB.CONTINUOUS, name="shake_end"
        )
        shake_use = self.model.addVars(
            self.num_workpieces, self.num_shake_eq,
            vtype=GRB.BINARY, name="shake_use"
        )

        # 分析阶段：每个工件i独立的开始/结束时间，及使用的设备e
        analyze_start = self.model.addVars(
            self.num_workpieces,  # 每个工件单独的开始时间
            vtype=GRB.CONTINUOUS, name="analyze_start"
        )
        analyze_end = self.model.addVars(
            self.num_workpieces,  # 每个工件单独的结束时间
            vtype=GRB.CONTINUOUS, name="analyze_end"
        )
        analyze_use = self.model.addVars(
            self.num_workpieces, self.num_analysis_eq,
            vtype=GRB.BINARY, name="analyze_use"  # 工件i是否使用设备e
        )

        # 最大完成时间（目标函数）
        makespan = self.model.addVar(vtype=GRB.CONTINUOUS, name="makespan")

        # -------------------------- 2. 约束条件 --------------------------
        # 2.1 清洗阶段约束
        # 每个工件仅用1台清洗设备
        for i in range(self.num_workpieces):
            self.model.addConstr(gp.quicksum(clean_use[i, c] for c in range(self.num_clean_eq)) == 1)
        # 清洗结束时间 = 开始时间 + 清洗时间（仅当使用该设备）
        for i in range(self.num_workpieces):
            for c in range(self.num_clean_eq):
                self.model.addConstr(clean_end[i, c] == clean_start[i, c] + self.clean_time * clean_use[i, c])
        # 同一清洗设备上的工件时间不重叠
        for c in range(self.num_clean_eq):
            for i in range(self.num_workpieces):
                for j in range(i + 1, self.num_workpieces):
                    # 若i和j都用设备c，则i的开始时间 >= j的结束时间 或 j的开始时间 >= i的结束时间
                    self.model.addConstr(
                        clean_start[i, c] >= clean_end[j, c] - (1 - clean_use[i, c]) * GRB.INFINITY
                    )
                    self.model.addConstr(
                        clean_start[j, c] >= clean_end[i, c] - (1 - clean_use[j, c]) * GRB.INFINITY
                    )

        # 2.2 充装阶段约束
        # 每个工件仅用1台充装设备
        for i in range(self.num_workpieces):
            self.model.addConstr(gp.quicksum(inflate_use[i, f] for f in range(self.num_inflation_eq)) == 1)
        # 充装结束时间 = 开始时间 + 充装时间（仅当使用该设备）
        for i in range(self.num_workpieces):
            for f in range(self.num_inflation_eq):
                self.model.addConstr(
                    inflate_end[i, f] == inflate_start[i, f] + self.inflation_time[i] * inflate_use[i, f])
        # 同一充装设备上的工件时间不重叠
        for f in range(self.num_inflation_eq):
            for i in range(self.num_workpieces):
                for j in range(i + 1, self.num_workpieces):
                    self.model.addConstr(
                        inflate_start[i, f] >= inflate_end[j, f] - (1 - inflate_use[i, f]) * GRB.INFINITY
                    )
                    self.model.addConstr(
                        inflate_start[j, f] >= inflate_end[i, f] - (1 - inflate_use[j, f]) * GRB.INFINITY
                    )
        # 充装需在清洗完成后开始（工件i的充装开始时间 >= 自身清洗结束时间）
        for i in range(self.num_workpieces):
            self.model.addConstr(
                gp.quicksum(inflate_start[i, f] * inflate_use[i, f] for f in range(self.num_inflation_eq))
                >= gp.quicksum(clean_end[i, c] * clean_use[i, c] for c in range(self.num_clean_eq))
            )

        # 2.3 摇匀阶段约束
        # 每个工件仅用1台摇匀设备
        for i in range(self.num_workpieces):
            self.model.addConstr(gp.quicksum(shake_use[i, s] for s in range(self.num_shake_eq)) == 1)
        # 摇匀时间 = 组分数 * 单位时间，结束时间 = 开始时间 + 摇匀时间
        shake_duration = self.num_components * self.shake_unit_time
        for i in range(self.num_workpieces):
            for s in range(self.num_shake_eq):
                self.model.addConstr(shake_end[i, s] == shake_start[i, s] + shake_duration * shake_use[i, s])
        # 同一摇匀设备上的工件时间不重叠
        for s in range(self.num_shake_eq):
            for i in range(self.num_workpieces):
                for j in range(i + 1, self.num_workpieces):
                    self.model.addConstr(
                        shake_start[i, s] >= shake_end[j, s] - (1 - shake_use[i, s]) * GRB.INFINITY
                    )
                    self.model.addConstr(
                        shake_start[j, s] >= shake_end[i, s] - (1 - shake_use[j, s]) * GRB.INFINITY
                    )
        # 摇匀需在充装完成后开始（工件i的摇匀开始时间 >= 自身充装结束时间）
        for i in range(self.num_workpieces):
            self.model.addConstr(
                gp.quicksum(shake_start[i, s] * shake_use[i, s] for s in range(self.num_shake_eq))
                >= gp.quicksum(inflate_end[i, f] * inflate_use[i, f] for f in range(self.num_inflation_eq))
            )

        # 2.4 分析阶段约束（核心修正）
        # 每个工件i使用的设备组合必须覆盖其所有组分
        for i in range(self.num_workpieces):
            for comp in range(self.num_components):  # 对每个组分comp
                # 至少有1台被使用的设备覆盖comp
                self.model.addConstr(
                    gp.quicksum(
                        analyze_use[i, e] for e in range(self.num_analysis_eq)
                        if comp in self.analysis_eq_capabilities[e]
                    ) >= 1
                )
        # 每个工件i至少使用1台分析设备
        for i in range(self.num_workpieces):
            self.model.addConstr(gp.quicksum(analyze_use[i, e] for e in range(self.num_analysis_eq)) >= 1)

        # 分析时间 = 工件i使用的所有设备中，对应组分的最长分析时间
        for i in range(self.num_workpieces):
            # 计算工件i的最大分析时间（所有使用的设备中，其覆盖组分的分析时间的最大值）
            max_analyze_time = self.model.addVar(vtype=GRB.CONTINUOUS, name=f"max_analyze_time_{i}")
            for e in range(self.num_analysis_eq):
                for comp in self.analysis_eq_capabilities[e]:
                    # 若使用设备e，则当前组分comp的时间可能影响最大值
                    self.model.addConstr(
                        max_analyze_time >= self.analysis_time[i, comp] * analyze_use[i, e]
                    )
            # 分析结束时间 = 开始时间 + 最大分析时间
            self.model.addConstr(analyze_end[i] == analyze_start[i] + max_analyze_time)

        # 分析需在摇匀完成后开始（工件i的分析开始时间 >= 自身摇匀结束时间）
        for i in range(self.num_workpieces):
            self.model.addConstr(
                analyze_start[i] >= gp.quicksum(shake_end[i, s] * shake_use[i, s] for s in range(self.num_shake_eq))
            )

        # 同一分析设备e上，工件的分析时间不重叠
        for e in range(self.num_analysis_eq):
            for i in range(self.num_workpieces):
                for j in range(i + 1, self.num_workpieces):
                    # 若i和j都用设备e，则i的开始时间 >= j的结束时间 或 j的开始时间 >= i的结束时间
                    self.model.addConstr(
                        analyze_start[i] >= analyze_end[j] - (2 - analyze_use[i, e] - analyze_use[j, e]) * GRB.INFINITY
                    )
                    self.model.addConstr(
                        analyze_start[j] >= analyze_end[i] - (2 - analyze_use[i, e] - analyze_use[j, e]) * GRB.INFINITY
                    )

        # 2.5 最大完成时间 = 所有工件分析结束时间的最大值
        self.model.addConstr(makespan >= gp.quicksum(analyze_end[i] for i in range(
            self.num_workpieces)) / self.num_workpieces * self.num_workpieces)  # 等效于max(analyze_end)
        for i in range(self.num_workpieces):
            self.model.addConstr(makespan >= analyze_end[i])

        # -------------------------- 3. 目标函数 --------------------------
        self.model.setObjective(makespan, GRB.MINIMIZE)

    def solve(self, time_limit: int = 300) -> Dict:
        self.model.setParam("TimeLimit", time_limit)
        self.model.optimize()

        if self.model.status == GRB.OPTIMAL:
            print(f"最优解: 最大完成时间 = {self.model.objVal:.2f}")
            return {
                "status": "optimal",
                "makespan": self.model.objVal,
                "variables": self._extract_variables()
            }
        elif self.model.status == GRB.TIME_LIMIT:
            print(f"超时，当前最佳解: 最大完成时间 = {self.model.objVal:.2f}")
            return {
                "status": "time_limit",
                "makespan": self.model.objVal,
                "variables": self._extract_variables()
            }
        else:
            print(f"求解失败，状态码: {self.model.status}")
            return {"status": "failed", "code": self.model.status}

    def _extract_variables(self) -> Dict:
        """提取关键变量结果（简化版）"""
        variables = {"clean": {}, "inflate": {}, "shake": {}, "analyze": {}}
        # 清洗阶段
        for i in range(self.num_workpieces):
            variables["clean"][i] = {
                c: {"start": self.model.getVarByName(f"clean_start[{i},{c}]").x,
                    "end": self.model.getVarByName(f"clean_end[{i},{c}]").x,
                    "used": self.model.getVarByName(f"clean_use[{i},{c}]").x}
                for c in range(self.num_clean_eq)
            }
        # 分析阶段
        for i in range(self.num_workpieces):
            variables["analyze"][i] = {
                "start": self.model.getVarByName(f"analyze_start[{i}]").x,
                "end": self.model.getVarByName(f"analyze_end[{i}]").x,
                "used_eq": [e for e in range(self.num_analysis_eq)
                            if self.model.getVarByName(f"analyze_use[{i},{e}]").x > 0.5]
            }
        return variables


if __name__ == "__main__":
    # 测试小型算例（容易求解）
    instance_path = "gas_scheduling_instances/small/instance_1.json"
    solver = GasSchedulingGurobiSolver(instance_path)
    solver.build_model()
    result = solver.solve(time_limit=600)  # 10分钟超时
    print(json.dumps(result, indent=2, ensure_ascii=False))